package com.zhz.mall.common.util.threadpool;

import cn.hutool.core.lang.Assert;
import com.alibaba.ttl.TtlCallable;
import com.alibaba.ttl.TtlRunnable;
import com.zhz.mall.common.config.BaseThreadPoolConfiguration;
import com.zhz.mall.common.util.context.SpringApplicationContextHolder;
import org.springframework.boot.autoconfigure.task.TaskExecutionAutoConfiguration;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;

import java.util.Date;
import java.util.Objects;
import java.util.concurrent.*;

/**
 * ThreadPoolUtils
 * <p>
 * 线程池工具类
 *
 * @author zhouhengzhe
 * @see BaseThreadPoolConfiguration
 */
public abstract class ThreadPoolUtils {

    private static final String APPLICATION_TASK_EXECUTOR_BEAN_NAME = TaskExecutionAutoConfiguration.APPLICATION_TASK_EXECUTOR_BEAN_NAME;
    private static final String APPLICATION_TASK_SCHEDULER_BEAN_NAME = "taskScheduler";
    private static ThreadPoolTaskExecutor applicationTaskExecutor;
    private static ThreadPoolTaskScheduler applicationTaskScheduler;


    private static ThreadPoolTaskExecutor getApplicationTaskExecutor() {
        if (Objects.isNull(applicationTaskExecutor)) {
            applicationTaskExecutor = SpringApplicationContextHolder.getApplicationContext().getBean(APPLICATION_TASK_EXECUTOR_BEAN_NAME, ThreadPoolTaskExecutor.class);
        }
        return applicationTaskExecutor;
    }

    public static ThreadPoolTaskExecutor getThreadPoolExecutor(String threadPoolName) {
        return SpringApplicationContextHolder.getApplicationContext().getBean(threadPoolName, ThreadPoolTaskExecutor.class);
    }

    private static ThreadPoolTaskScheduler getApplicationTaskScheduler() {
        if (Objects.isNull(applicationTaskScheduler)) {
            applicationTaskScheduler = SpringApplicationContextHolder.getApplicationContext().getBean(APPLICATION_TASK_SCHEDULER_BEAN_NAME, ThreadPoolTaskScheduler.class);
        }
        return applicationTaskScheduler;
    }

    public static void execute(Runnable task) {
        getApplicationTaskExecutor().execute(task);
    }

    public static Future<?> submit(Runnable task) {
        return getApplicationTaskExecutor().submit(task);
    }

    public static <V> Future<V> submit(Callable<V> task) {
        return getApplicationTaskExecutor().submit(task);
    }

    public static ScheduledFuture<?> schedule(Runnable task, Date startTime) {
        return getApplicationTaskScheduler().schedule(task, startTime);
    }

    public static ScheduledFuture<?> schedule(Runnable task, int delayTimeMillis) {
        Date startTime = new Date(System.currentTimeMillis() + delayTimeMillis);
        return schedule(task, startTime);
    }

    /**
     * 执行任务（可为目标任务单独指定拒绝策略）
     */
    public static void execute(Runnable task, RejectedExecutionHandler rejectedExecutionHandler) {
        Assert.notNull(rejectedExecutionHandler);
        try {
            execute(task);
        } catch (RejectedExecutionException e) {
            rejectedExecutionHandler.rejectedExecution(task, getApplicationTaskExecutor().getThreadPoolExecutor());
        }
    }

    /**
     * 执行任务（当线程池拒绝执行时是否由当前线程执行）
     */
    public static void execute(Runnable task, boolean callerRunsPolicy) {
        if (callerRunsPolicy) {
            execute(task, new ThreadPoolExecutor.CallerRunsPolicy());
        } else {
            execute(task);
        }
    }


    /**
     * 等待 future 执行完毕, 若执行异常抛出实际产生的异常, 而不是 ExecutionException
     */
    public static <V> V get(Future<V> future) throws Throwable {
        try {
            return future.get();
        } catch (ExecutionException e) {
            if (Objects.nonNull(e.getCause())) {
                throw e.getCause();
            }
            throw e;
        }
    }


    /**
     * 等待 future 执行完毕, 若执行异常抛出实际产生的异常, 而不是 ExecutionException
     */
    public static <V> V get(Future<V> future, long timeout, TimeUnit timeUnit) throws Throwable {
        try {
            return future.get(timeout, timeUnit);
        } catch (ExecutionException e) {
            if (Objects.nonNull(e.getCause())) {
                throw e.getCause();
            }
            throw e;
        }
    }


    /**
     * 若使用了 TransmittableThreadLocal 则提交任务时应使用 TtlRunnable/TtlCallable 进行包装
     */
    public static <V> Future<V> submitWithTtl(Callable<V> task) {
        task = TtlCallable.get(task, false, true);
        return getApplicationTaskExecutor().submit(Objects.requireNonNull(task));
    }


    public static void executeWithTtl(Runnable task) {
        task = TtlRunnable.get(task, false, true);
        getApplicationTaskExecutor().execute(Objects.requireNonNull(task));
    }



    public static ScheduledFuture<?>  scheduleWithTtl(Runnable task, int delayTimeMillis) {
        task = TtlRunnable.get(task, false, true);
        return schedule(Objects.requireNonNull(task), delayTimeMillis);
    }

    /**
     * 等待 future 执行完毕（忽略执行异常）
     *
     * @param future              task future
     * @param timeout             超时时间
     * @param timeUnit            超时时间单位
     * @param autoCancelIfTimeOut 当超时时是否自动取消任务
     * @return 异常信息
     */
    public static Throwable safeWait(Future<?> future, long timeout, TimeUnit timeUnit, boolean autoCancelIfTimeOut) {
        Throwable cause = null;

        try {
            get(future, timeout, timeUnit);
        } catch (Throwable t) {
            cause = t;
        }

        if (autoCancelIfTimeOut && Objects.nonNull(cause) && cause instanceof TimeoutException) {
            future.cancel(true);
        }

        return cause;
    }


    /**
     * 等待 future 执行完毕（忽略执行异常）
     */
    public static Throwable safeWait(Future<?> future, long timeout, TimeUnit timeUnit) {
        return safeWait(future, timeout, timeUnit, false);
    }


}
